/*
 *
 * Copyright (C) 2008 Ed Hong
 * 
 * This program is free software; you can redistribute it and/or
 * modify it under the terms of the GNU General Public License
 * as published by the Free Software Foundation; either version 2
 * of the License, or (at your option) any later version.
 * 
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 * 
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.
 *
 */

#include "contexttree.h"




static int find_context_num_rec(
		     const SubbandComponent *sc, int componentnum, // component num for color images
		     int r, int c, int sb, int bp, const NeighborSignificance *ns1, 
		     const NeighborSigns *ns2,
		     const contexttreenode *ct) {
  if (ct->leafnum != -1) return ct->leafnum; 
  else {
    Characteristic *ch;
    ch = &AllC[ct->charindex];
    if (ch->f == NULL) {
      printf("find_context_num_rec: function pointer for Characteristic %s does not exist; exitting..\n", ch->name);
      exit(1);
    }
    else {
      int bitcharval;
      contexttreenode *newct;
      bitcharval = (*(ch->f))(sc, componentnum, ns1, ns2, r, c, sb, bp);
      if (bitcharval <= ct->charvalue) {
	newct = ct->left;
      }
      else  {
	newct = ct->right;
      }
      return find_context_num_rec(sc, componentnum, r, c, sb, bp, ns1, ns2, newct);
    }

  }

}



int find_context_num(
		     const SubbandComponent *sc, int componentnum, // component num for color images
		     int r, int c, int sb, int bp, const contexttree *t) {
  NeighborSignificance nsig; 
  NeighborSigns nsign;

  computeNeighborSignificance(&sc[componentnum], r,c,sb,bp, &nsig);
  computeNeighborSigns(&sc[componentnum], r,c,sb, bp, &nsign);

  return
    find_context_num_rec(sc, componentnum, r, c, sb, bp, &nsig, &nsign, t->root);

}


static void init_ContextIDMap(const unsigned long long *charset,  contexttree *x) {
  int *num;
  num= &(x->num_on_disk_chars);
  *num =0;
  int i;
  //  printf("charset%llX %llX\n", charset[1], charset[0]);
  for (i=0; i < TOTAL_NUM_CHARACTERISTICS; i++) {
    if (checkflag(i,charset)) {
      x->ContextIDMap[*num] = i;
      x->reverseContextIDMap[i] = *num;
      (*num)++;
    }
  }

}


int create_fullsplitcontexttree_rec(const unsigned long long *charset, contexttreenode *out, int index, int charvalindex, int *currleafnum){

  //  printf("createfulrec: %d, %d, %d\n", index, charvalindex, *currleafnum);
  int i;
  Characteristic *ch;
  i = index;
  while (i < TOTAL_NUM_CHARACTERISTICS){
    if (checkflag(i, charset)) {
      ch = &AllC[i];
      if (ch->f == NULL) {
	printf("create_fullsplitcontexttree_rec: function pointer for Characteristic %s does not exist; exitting..\n", ch->name);
	exit(1);
      }
      else {
	out->charindex = i;
	out->charvalue = charvalindex +  ch->valuestart;
	if (charvalindex +2 >= 	ch->numvalues) {
	  // last split at this index
	  charvalindex = 0;
	  i++;
	}
	else charvalindex++;
	out->left = malloc (sizeof(contexttreenode));
	out->right = malloc (sizeof( contexttreenode));
	if (out->left == NULL || out->right == NULL) {
	  printf("create_fullsplitcontexttree_rec: error allocating memory exitting..\n");
	  return 1;
	}
	out->leafnum = -1;
	return create_fullsplitcontexttree_rec(charset, out->left, out->charindex+1, 0, currleafnum) + 
	  create_fullsplitcontexttree_rec(charset, out->right, i, charvalindex, currleafnum); 
	
      }      
    }
    i++;
  }
  out->charindex = -1;
  out->charvalue = -327; // arbitrary dummy value
  out->left = NULL;
  out->right = NULL;
  out->leafnum = *currleafnum;
  (*currleafnum)++;
  return 0;
}

static contexttreenode *create_contexttreenode(int chari, int charv, int leafn) {
  contexttreenode *newnode; 
  newnode =  malloc (sizeof( contexttreenode));
  if (newnode == NULL) {
	  printf("create_contexttreenode: error allocating memory exitting..\n");
	  exit(1);
  }
  newnode->charindex = chari;
  newnode->charvalue = charv;
  newnode->left = NULL;
  newnode->right =NULL;
  newnode->leafnum = leafn;
  newnode->charranges = NULL;
  return newnode;
}

int create_fullsplitcontexttree(const unsigned long long *charset, contexttree *out){ // tree that splits on 
// all context characteristics in charset.  
  int currleafnum;
  out->root =  malloc (sizeof( contexttreenode));
  if (out->root == NULL) {
	  printf("create_fullsplitcontexttree: error allocating memory exitting..\n");
	  return 1;
  }
  currleafnum=0;
  int result = create_fullsplitcontexttree_rec(charset, out->root, 0, 0, &currleafnum); 
  out->numleafnodes = currleafnum;
  
  init_ContextIDMap(charset, out);
  
  return result;
}


int create_imgsplitcontexttree(const unsigned long long *charset, contexttree *out){ // tree that splits on 
// all context characteristics in charset; refine, sig and sign characteristics special.   
  int currleafnum;
  unsigned long long newcharset[2];
  contexttreenode *sigroot, *issign, *isrefine, *issig;

  init_ContextIDMap(charset, out);

  issign = isrefine = issig = NULL;

  newcharset[0] = charset[0];
  newcharset[1] = charset[1];

  if (checkflag(IS_SIGN, newcharset)) {
    removeflag(IS_SIGN, newcharset);
    issign = create_contexttreenode(IS_SIGN,0,-1);
  }
  if (checkflag(IS_SIG, newcharset)) {
    removeflag(IS_SIG, newcharset);
    issig = create_contexttreenode(IS_SIG,0,-1);
  }
  if (checkflag(IS_REFINE, newcharset)) {
    removeflag(IS_REFINE, newcharset);
    isrefine = create_contexttreenode(IS_REFINE,0,-1);
  }
  sigroot =  malloc (sizeof( contexttreenode));
  if (sigroot == NULL) {
	  printf("create_imgsplitcontexttree: error allocating memory exitting..\n");
	  return 1;
  }
  int result;
  currleafnum = 0;

  // order is issig, isrefine, then issign on top. 

  if (issig) {
    out->root = issig;
    if ((isrefine!= NULL) || (issign != NULL) ) {
      if (isrefine) {
	out->root->left = isrefine;
	if (issign) {
	  free(issign);
	}
      }
      else {// isrefine == NULL, issign != NULL
	out->root->left = issign;
      }
      out->root->left->left = create_contexttreenode(-1, -1, currleafnum); 
      currleafnum++;
      out->root->left->right = create_contexttreenode(-1, -1, currleafnum); 
      currleafnum++;
    }
    else { // only is_sig on top
      out->root->left = create_contexttreenode(-1, -1, currleafnum); 
      currleafnum++;
    }    
    result = create_fullsplitcontexttree_rec(newcharset, sigroot, 0, 0, &currleafnum); 
    out->root->right = sigroot;
  }
  else if (isrefine) { // isrefine on top, no issig. 
    out->root = isrefine;
    result = create_fullsplitcontexttree_rec(newcharset, sigroot, 0, 0, &currleafnum); 
    if (issign) {
      out->root->left = issign;
      out->root->left->left = sigroot;
      out->root->left->right = create_contexttreenode(-1, -1, currleafnum); 
      currleafnum++;
    }
    else { // just single isrefine on top
      out->root->left = sigroot;
    }
    out->root->right = create_contexttreenode(-1, -1, currleafnum); 
    currleafnum++;
  }
  else if (issign) { // single issign on top.
    out->root = issign;
    result = create_fullsplitcontexttree_rec(newcharset, sigroot, 0, 0, &currleafnum); 
    out->root->left = sigroot;
    out->root->right  = create_contexttreenode(-1, -1, currleafnum); 
    currleafnum++;
  }
  else { // issig, isrefine, and issign all false. 
    result = create_fullsplitcontexttree_rec(newcharset, sigroot, 0, 0, &currleafnum); 
    out->root = sigroot;
  }

  out->numleafnodes = currleafnum;
  return result;
}

int getsplitnummax(const contexttreenode *n,
                    const contexttree *x);

// 80 spaces.
char SPACES[81] = "                                                                                ";
static void printContextTreeNodeNonrec(int numspaces, const contexttreenode *ct, const contexttree *x) {
  char *spacestring; 
  spacestring = SPACES + (80 - numspaces);
  printf("%saddress:%8X\n", spacestring, (int)ct);
  printf("%scharindex:%d\n", spacestring, ct->charindex);
  printf("%scharvalue:%d\n", spacestring, ct->charvalue);
  printf("%sleft: %8X\n", spacestring, (int)ct->left);
  printf("%sright:%8X\n", spacestring, (int)ct->right);
  printf("%sleafnum:%d\n", spacestring, ct->leafnum);
  if (ct->charranges !=NULL) {
    int i;
    int num_ondisk_chars = x->num_on_disk_chars;
    printf("%scharranges:", spacestring);
    for (i=0; i < num_ondisk_chars; i++) {
      if (ct->charranges[i].min == ct->charranges[i].max)  
        printf("(%d)", ct->charranges[i].min);
      else 
        printf("(%d,%d)", ct->charranges[i].min, ct->charranges[i].max);
    }
    printf ("\n");
    printf("%smaxsplitnum: %d\n", spacestring, getsplitnummax(ct, x));
  }
}
    
static void printcontexttreeNoderec(int numspaces, const contexttreenode *ct, const contexttree *x) {
    printContextTreeNodeNonrec(numspaces,ct, x);  
    if (ct->leafnum == -1) {
      printcontexttreeNoderec(numspaces+2,ct->left, x);
      printcontexttreeNoderec(numspaces+2,ct->right, x);
    }
}

void printcontexttree(const contexttree *x) {
   int i;
   printf("numleafnodes: %d\n",x->numleafnodes); 
 
   printf("num_on_disk_chars: %d\n", x->num_on_disk_chars);
   
   printf("ContextIDMap:");   
   for (i=0; i < x->num_on_disk_chars; i++) {
      printf(" %d", x->ContextIDMap[i]);
    }
   printf("\n");

   printf("reverseContextIDMap:");   
   for (i=0; i < TOTAL_NUM_CHARACTERISTICS; i++) {
      printf(" %d", x->reverseContextIDMap[i]);
    }
   printf("\n");
 
   printcontexttreeNoderec(0, x->root, x);
}


static contexttreenode *readctd_rec(FILE * f, int ContextIDMap[],  int *currleafnum) {
  contexttreenode *curr;
  int label, i1, i2;
  fscanf(f, "%d %d %d", &label, &i1, &i2);
  if (i1 == -1 && i2 ==-1) {
    // leaf node
    curr = create_contexttreenode(-1,-1,*currleafnum);

    (*currleafnum)++;

  }
  else {
    curr = create_contexttreenode(ContextIDMap[i1], i2, -1);
    //    printf("%d %d\n", curr->charindex, curr->charvalue);
    curr->left = readctd_rec(f, ContextIDMap, currleafnum);
    curr->right = readctd_rec(f, ContextIDMap, currleafnum);
  }
  return curr;
}



int readctdfile(const char *filename, const unsigned long long *charset, contexttree *out) {
  int currleafnum;

  init_ContextIDMap(charset, out);

  FILE *file;
  file = fopen(filename, "r");
  if (file == NULL) {
    QccErrorAddMessage("readctdfile: Error opening file %s", filename
			 );
    return 1;
  }

  currleafnum = 0;
  out->root = readctd_rec(file, out->ContextIDMap, &currleafnum);

  out->numleafnodes = currleafnum;
  return 0;
  
}

static void   printctdfile_rec(FILE *file, const contexttreenode *n, const int reverseContextIDMap[]) {
  int label = 0;
  if (n->leafnum != -1) {
    fprintf(file, "%d -1 -1\n", label);
  }
  else {
    fprintf(file, "%d %d %d\n", label, reverseContextIDMap[n->charindex], n->charvalue);
    printctdfile_rec(file, n->left, reverseContextIDMap);
    printctdfile_rec(file, n->right, reverseContextIDMap);
  }
}


int printctdfile(const char *filename, const contexttree *out) {
  FILE *file;

  file = fopen(filename, "w");
  printctdfile_rec(file, out->root, out->reverseContextIDMap);
  return 0;
}

// allocates and returns pointer to a new charrange;
// returns NULL on error;
static range *allocate_charrange(const contexttree *x) {
   return (range *) malloc(sizeof(range)*x->num_on_disk_chars);
}

// allocates (if necssary) and sets charrange of root
// returns 1 on error; 0 on success. 
static int init_root_charrange(const contexttree *x) {
    if (x->root->charranges != NULL) {
       return 0;
    }
    x->root->charranges = allocate_charrange(x);
    if (x->root->charranges == NULL) {
        QccErrorAddMessage("init_root_charrange: unable to allocate memory");
              return 1;
    }
    // initialize charrange to max possible ranges...
    int i;
    for (i=0; i < x->num_on_disk_chars; i++) {
       x->root->charranges[i].min = AllC[x->ContextIDMap[i]].valuestart;
       x->root->charranges[i].max = AllC[x->ContextIDMap[i]].valuestart + 
                                  AllC[x->ContextIDMap[i]].numvalues- 1;
    }

    return 0;
}

// sets charrange for the two children. 
// allocates charrange for the two children, if necessary. 
// assumes children are already allocated
// assumes charindex and charvalue for input contexttreenode
//   correctly set. 
// returns 1 on allocate failure;
// returns 2 on charindex not set correctly.
// returns 0 on success.  
static int set_children_charrange(contexttreenode *ct, const contexttree *x) {
   if (ct->charindex < 0 || ct->charindex >= TOTAL_NUM_CHARACTERISTICS) {
      QccErrorAddMessage("set_children_charrange: bad charindex of %d", ct->charindex);
      return 2;
   }
   if (ct->left->charranges == NULL) {
      ct->left->charranges = allocate_charrange(x);
   }
   if (ct->right->charranges == NULL) {
      ct->right->charranges = allocate_charrange(x);
   }

   if (ct->left->charranges == NULL || ct->right->charranges == NULL) {
      QccErrorAddMessage("set_children_charrange: unable to allocate memory");
      return 1;
   }
   
   memcpy(ct->left->charranges, ct->charranges,sizeof(range)*x->num_on_disk_chars );
   memcpy(ct->right->charranges, ct->charranges,sizeof(range)*x->num_on_disk_chars );
   int ondiskcharid = x->reverseContextIDMap[ct->charindex];
   ct->left->charranges[ondiskcharid].max = ct->charvalue;
   ct->right->charranges[ondiskcharid].min = ct->charvalue+1;
   return 0;
}

static int calc_num_bits(int numpossiblevalues) {
   int numbits = 0;
    int maxvalue = numpossiblevalues - 1;
    while (maxvalue > 0) {
        numbits++;
        maxvalue = maxvalue >> 1;
    }
    return numbits;
}

// encode value using least number of bits necessary. 
//  numpossiblevalues is  1 + largest value possible for value; 
//  minimum possible value assumed to be zero. 
//  return 1 with invalid initial input or error encoding value;
//  return 0 = no error. 
// if numpossiblevalues = 1, will return without error and 
//   without outputting any bits into output_buffer. 
static int bitbuffer_encode_num(QccBitBuffer *output_buffer, int value, int numpossiblevalues) {
    if (numpossiblevalues <=0)  {
        printf("bitbuffer_encode_num: cannot have numpossiblevalues of %d\n", 
                numpossiblevalues);
        return 1;
    }  
    if (value < 0 || (value >= numpossiblevalues)) {
        printf("bitbuffer_encode_num: value %d invalid when numpossiblevalues= %d\n", 
                value, numpossiblevalues);
        return 1;    
    }
    int numbits = calc_num_bits(numpossiblevalues);
     if (numbits == 0) {
        // no bits necessary, don't encode. 
        return 0;
    }
    else {
//        printf("encoding a %d using %d bits (%d possible)\n", value, numbits,numpossiblevalues);
        return QccBitBufferPutBits(output_buffer, value, numbits);
    }
}

// decode a number  using least number of bits necessary, storing
//   result in value.  
//  numpossiblevalues is  1 + largest value possible for value; 
//  minimum possible value assumed to be zero. 
//  return 1 with invalid initial input or error encoding value;
//  return 0 = no error. 
// if numpossiblevalues = 1, will return without error and 
//   without reading any bits from input_buffer, with *value set to 0.   
static int bitbuffer_decode_num(QccBitBuffer *input_buffer, int *value, int numpossiblevalues) {
    if (numpossiblevalues <=0)  {
        printf("bitbuffer_decode_num: cannot have numpossiblevalues of %d\n", 
                numpossiblevalues);
        return 1;
    }  
    
    int numbits = calc_num_bits(numpossiblevalues);

    if (numbits == 0) {
        // no bits necessary, don't read any bits. 
        *value = 0;
        return 0;
    }
    else {
        int error = QccBitBufferGetBits(input_buffer, value, numbits);
//        printf("decoded a %d with %d bits (%d possible)\n", *value, numbits, numpossiblevalues);
        return error;
    }
}

// get max # possible splits for conttexttree node; assumes charranges set.
int getsplitnummax(const contexttreenode *n,
                    const contexttree *x) {
     int currsplitnum; 
     int i;
     currsplitnum = 0;
     for (i=0; i < x->num_on_disk_chars; i++) {
        int nextCharIndex;
        nextCharIndex = x->ContextIDMap[i];
         currsplitnum += n->charranges[i].max - n->charranges[i].min; 
     }
     return currsplitnum;
}                    


// returns splitnum based on charindex and charvalue fields of n.
// quits and exits on failure. 
int getsplitidnum( const contexttreenode *n,
                    const contexttree *x) {
     int currsplitnum; 
     int i;
     currsplitnum = 0;
     for (i=0; i < x->num_on_disk_chars; i++) {
        int nextCharIndex;
        nextCharIndex = x->ContextIDMap[i];
        if (n->charindex == nextCharIndex) {
           return currsplitnum + n->charvalue -  n->charranges[i].min;
        }
        currsplitnum += n->charranges[i].max - n->charranges[i].min; 
     }
     printf("getsplitidnum: error getting id of split\n");
     exit(1);
     return -1;
}                    

// returns 1 on error, 0 on sucess
// sets charid and charval based on splitidnum
int getcharidandcharvalue(int *charid, int *charval, int splitidnum, const contexttreenode *n, 
                            const contexttree *x) {
     int currsplitnum; 
     int i;
     currsplitnum = 0;
     for (i=0; i < x->num_on_disk_chars; i++) {
        int nextCharIndex;
        nextCharIndex = x->ContextIDMap[i];
        if (splitidnum >= currsplitnum && 
             splitidnum < currsplitnum + n->charranges[i].max - n->charranges[i].min) {
           *charid = nextCharIndex;
           *charval = splitidnum - currsplitnum + n->charranges[i].min;
           return 0;
        }
        currsplitnum += n->charranges[i].max - n->charranges[i].min;
     }
     QccErrorAddMessage("getcharidnadcharvalue: error finding char id  and value of split id %d", 
                       splitidnum );
     return 1;
     
}

// assumes n->charranges already created and set appopriately before call. 
// also assumes charranges set correctly. 
static int encodectd_bitbuffer_rec(QccBitBuffer *output_buffer, const contexttreenode *n, 
                                    const contexttree *x) {
  if (n->leafnum != -1) {
    return bitbuffer_encode_num(output_buffer, 1, 2);
  }
  else {
    int error, splitnum, splitnummax; 
    error = bitbuffer_encode_num(output_buffer, 0, 2);

    splitnum = getsplitidnum(n,x);
    splitnummax = getsplitnummax(n,x);

    error += bitbuffer_encode_num(output_buffer, splitnum, splitnummax); 

    error += encodectd_bitbuffer_rec(output_buffer, n->left, x);
    error += encodectd_bitbuffer_rec(output_buffer, n->right, x);
    return error; 
  }
}

static int set_charrange_rec(contexttreenode *ct, const contexttree *x) {
   if (ct->leafnum == -1) {
     int errors = 0;
     errors = (set_children_charrange(ct, x) != 0);
     errors += set_charrange_rec(ct->left, x);
     errors += set_charrange_rec(ct->right, x);
     return errors;
   }
   else return 0;
}

// Encodes nodes in an pre-order traversal. 
//  each node represented by:
//      1) single bit, 1 for is_leafnode, 0 for not leaf node
//      if this is not a leaf node, also use
//         2)  a splitnum ID to represent the split. 
//         can calculate total num of possible splits through
//         computing appropriate charranges in each node. 
// sets charranges before starting the encode process.
// returns # of errors that occurred.     
int encodectd_bitbuffer(QccBitBuffer *output_buffer, const contexttree *out) {
  int numerrors = 0;
  numerrors = init_root_charrange(out);
  numerrors += set_charrange_rec(out->root, out);
  numerrors += encodectd_bitbuffer_rec(output_buffer, out->root, out);
  return numerrors;
}

// assumes curr already allocated (to dummy initial values) 
// returns # errors; 0 on success. 
static int decodectd_bitbuffer_rec(QccBitBuffer *input_buffer,
                        contexttreenode *curr, 
                        contexttree *x,  int *currleafnum) { 

  int  splitidnum, isleafnode;
  int errorval; 

  splitidnum = -1; // dummy intialization
  isleafnode = -1; // dummy initialization.

  errorval = 0;
  errorval += bitbuffer_decode_num(input_buffer, &isleafnode, 2);
  
  if (isleafnode) {
    curr->leafnum = *currleafnum;
    (*currleafnum)++;    
  }
  else {
    int splitnummax;
    splitnummax = getsplitnummax(curr,x);
    errorval += bitbuffer_decode_num(input_buffer, &splitidnum, splitnummax);

    errorval += getcharidandcharvalue(&(curr->charindex), &(curr->charvalue), splitidnum, curr, x);
      
    curr->left = create_contexttreenode(-1, -1, -1);
    curr->right = create_contexttreenode(-1, -1, -1);
    errorval += (set_children_charrange(curr,x)!= 0);

    errorval += decodectd_bitbuffer_rec(input_buffer, curr->left, x,
                            currleafnum);
    errorval += decodectd_bitbuffer_rec(input_buffer, curr->right, x,
                            currleafnum);
  }
  return errorval;
}


// return value represents whether error occurred. 
// decodes ctd tree from input_buffer; 
//  see encode for format of ctd tree. 
//   returns # of errors that occurred.     
int decodectd_bitbuffer(QccBitBuffer *input_buffer, const unsigned long long *charset,
                        contexttree *out) {
  int  currleafnum, errors;
  init_ContextIDMap(charset, out );

  currleafnum = 0;
  out->root = create_contexttreenode(-1,-1,-1);
  errors = 0;
  errors += init_root_charrange(out);
  errors += decodectd_bitbuffer_rec(input_buffer, out->root, out, 
                        &currleafnum);

  out->numleafnodes = currleafnum;
  return errors;
}


static void rec_printtree(const contexttreenode *n, const int reverseContextIDMap[]) {
  printf("(\n");
  if (n->leafnum != -1) {
    printf(" () ");
  }
  else {
    printf( " %d %d ",  reverseContextIDMap[n->charindex], n->charvalue);
    rec_printtree(n->left, reverseContextIDMap);
    rec_printtree(n->right, reverseContextIDMap);
  }
  
  printf(")\n");
}

void printtree(const contexttree *x) {
  printf("\\font\\MySmall=cmr10 at 5pt\n \\MySmall \n \\hsize=100in \\vsize=20in\n");

  printf("\\tree\n");
  rec_printtree(x->root, x->reverseContextIDMap);
  printf("\n\\bye\n");
}

